#ifndef MAXPOOL_2X2_H_
#define MAXPOOL_2X2_H_

#include "stream_tools.h"

template<unsigned BIT>
ap_uint<BIT> max(const ap_uint<BIT>& x, const ap_uint<BIT>& y)
{
    return x > y ? x : y;
};

template<unsigned N_IO, unsigned BIT>
ap_uint<N_IO * BIT> max(const ap_uint<N_IO * BIT>& x, const ap_uint<N_IO * BIT>& y)
{
	#pragma HLS INLINE
    ap_uint<N_IO * BIT> z;
    for (unsigned i = 0; i < N_IO; ++i)
    {
        #pragma HLS UNROLL
        ap_uint<BIT> a = x(SLICE(BIT, i));
        ap_uint<BIT> b = y(SLICE(BIT, i));
        ap_uint<BIT> c = a > b ? a : b;
        z(SLICE(BIT, i)) = c;
    }
    return z;
};


template<
    unsigned N_CH,
    unsigned N_IH,
    unsigned N_IW,
    unsigned N_BATCH,
    unsigned BIT_ACTV
>
void maxpool2x2_golden(data_stream<BIT_ACTV>& in, data_stream<BIT_ACTV>& out) {
    constexpr unsigned N_OH = N_IH / 2;
    constexpr unsigned N_OW = N_IW / 2;

    ap_uint<BIT_ACTV> input_buf[N_BATCH][N_IH][N_IW][N_CH];
    ap_uint<BIT_ACTV> output_buf[N_BATCH][N_OH][N_OW][N_CH];

    // Read input data
    for(unsigned b = 0; b < N_BATCH; ++b) {
        for(unsigned ih = 0; ih < N_IH; ++ih) {
            for(unsigned iw = 0; iw < N_IW; ++iw) {
                for(unsigned ic = 0; ic < N_CH; ++ic) {
                    input_buf[b][ih][iw][ic] = in.read();
                }
            }
        }
    }
    
    // Perform maxpool operation
    for(unsigned b = 0; b < N_BATCH; ++b) {
        for(unsigned oh = 0; oh < N_OH; ++oh) {
            for(unsigned ow = 0; ow < N_OW; ++ow) {
                for(unsigned ic = 0; ic < N_CH; ++ic) {
                    ap_uint<BIT_ACTV> val00 = input_buf[b][oh * 2][ow * 2][ic];
                    ap_uint<BIT_ACTV> val01 = input_buf[b][oh * 2][ow * 2 + 1][ic];
                    ap_uint<BIT_ACTV> val10 = input_buf[b][oh * 2 + 1][ow * 2][ic];
                    ap_uint<BIT_ACTV> val11 = input_buf[b][oh * 2 + 1][ow * 2 + 1][ic];

                    ap_uint<BIT_ACTV> max01 = max<BIT_ACTV>(val00, val01);
                    ap_uint<BIT_ACTV> max23 = max<BIT_ACTV>(val10, val11);
                    ap_uint<BIT_ACTV> result = max<BIT_ACTV>(max01, max23);
                    output_buf[b][oh][ow][ic] = result;
                }
            }
        }
    }

    // Write output data
    for(unsigned b = 0; b < N_BATCH; ++b) {
        for(unsigned oh = 0; oh < N_OH; ++oh) {
            for(unsigned ow = 0; ow < N_OW; ++ow) {
                for(unsigned ic = 0; ic < N_CH; ++ic) {
                    out.write(output_buf[b][oh][ow][ic]);
                }
            }
        }
    }
    
    assert(in.empty());
    assert(out.size() == N_BATCH * N_OH * N_OW * N_CH);
    return;
};

template<
    unsigned P_CH,
    unsigned N_CH,
    unsigned N_IH,
    unsigned N_IW,
    unsigned N_BATCH,
    unsigned BIT_ACTV
>
void maxpool2x2(data_stream<P_CH * BIT_ACTV>& in, data_stream<P_CH * BIT_ACTV>& out) {
    static_assert(N_CH >= P_CH, "maxpool2x2");
    static_assert(N_CH % P_CH == 0, "maxpool2x2");

    constexpr unsigned FOLD = N_CH / P_CH;
    constexpr unsigned ITER = N_BATCH * N_IH * N_IW * FOLD;
    assert(in.size() == ITER);
    assert(out.empty());

    #pragma HLS DATAFLOW

    ap_uint<P_CH * BIT_ACTV> line_buf[N_IW / 2][FOLD];

    for (unsigned r = 0; r < N_BATCH * N_IH; ++r)
    {
        for (unsigned c = 0; c < N_IW; ++c)
        {
            for (unsigned f = 0; f < FOLD; ++f)
            {
                #pragma HLS PIPELINE II=1
                const unsigned idx = c >> 1;
                ap_uint<P_CH * BIT_ACTV> in_buf = in.read();
                ap_uint<P_CH * BIT_ACTV> out_buf;
                if ((r & 0x1) == 0) {
                    if ((c & 0x1) == 0) {
                        // 0x0
                        line_buf[idx][f] = in_buf;
                    } else {
                        // 0x1
                        out_buf = max<P_CH, BIT_ACTV>(in_buf, line_buf[idx][f]);
                        line_buf[idx][f] = out_buf;
                    }
                } else {
                    if ((c & 0x1) == 0) {
                        // 0x2
                        out_buf = max<P_CH, BIT_ACTV>(in_buf, line_buf[idx][f]);
                        line_buf[idx][f] = out_buf;
                    } else {
                        // 0x3
                        out_buf = max<P_CH, BIT_ACTV>(in_buf, line_buf[idx][f]);
                        out.write(out_buf);
                    }
                }
            }
        }
    }

    assert(in.empty());
    assert(out.size() == ITER / 4);
    return;
};

#endif